import torch
import numpy
import torch.nn as nn
import torch.nn.functional as F
import math
import os
from PIL import Image
import datetime
from torch.utils.data import DataLoader
from dataset import CustomDataset
from models import DFN
from utils import Lossor, Evaluator



def save_results(preds, num_class, result_dir, image_name):
    COLOR_MAP = [(128, 64, 128), (244, 35, 232), (70, 70, 70), (102, 102, 156), (190, 153, 153), (153, 153, 153),
                 (250, 170, 30), (220, 220, 0), (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60),
                 (255, 0, 0), (0, 0, 142), (0, 0, 70), (0, 60, 100), (0, 80, 100), (0, 0, 230), (119, 11, 32)]
    for i in range(len(preds)):
        pre = preds[i]
        row, col = pre.shape
        dst = torch.zeros((row, col, 3))
        for j in range(num_class):
            dst[pre == j] = torch.Tensor(COLOR_MAP[j])
        dst = numpy.array(dst, dtype=numpy.uint8)
        dst = Image.fromarray(dst)
        dst = dst.resize((dst.size[0]*4, dst.size[1]*4), Image.NEAREST)
        dst.save(result_dir+'/'+image_name[i], 'PNG')

def main():
    batch_size = 3
    start_cuda = 'cuda:0'
    gpu_ids = [0]

    num_class = 19
    num_show = 2
    ckpt_path = './ckpt'
    result = './result'
    save_result = True

    encoder = 'resnet' #encoder = 'shufflenet'#
    net_weight_path = 'resnet-dfn-v1/model_epoch_99.pth' #net_weight_path = 'shufflenet-dfn-v1/model_epoch_5.pth' #

    dataset=CustomDataset('leftImg8bit', 'gtFine', '_leftImg8bit.png', '_gtFine_labelTrainIds.png', task='val')
    evalloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    # focal_weight = 0.1
    model_id = encoder+'-dfn-v1'
    result_dir = result+'/'+model_id+datetime.datetime.now().strftime('[%Y%m%d%H%M]')
    if not os.path.isdir(result_dir):
        os.makedirs(result_dir)

    net = DFN(num_class, weights=None, encoder=encoder)


    device = torch.device(start_cuda if torch.cuda.is_available() else "cpu")
    if torch.cuda.device_count() > 1:
        net = nn.DataParallel(net, device_ids=gpu_ids)
    if net_weight_path:
        print('loading the net weight : {}'.format(net_weight_path))
        net.load_state_dict(torch.load(ckpt_path + '/' + net_weight_path))
        net.eval()
    net.to(device)

    # In[ ]:

    #training
    print('start evaluate on GPU : {}'.format(gpu_ids))
    # crit = Lossor(focal_weight)
    evaluator = Evaluator(num_class)

    for i, data in enumerate(evalloader, 0):
        inputs, labels, image_name = data
        inputs = inputs.float()
        #         labels = labels.float()
        inputs = inputs.to(device)
        labels = labels.to(device)
        outputs = net(inputs)
        # loss = crit.MixLoss(outputs, labels)
        preds = torch.argmax(outputs[-2].data, 1)
        evaluator.add_batch(labels, preds)
        if save_result:
            save_results(preds, num_class, result_dir, image_name)
        if i % num_show == num_show - 1:
            # loss_tmp = crit.mean_loss()
            mIoU_tmp = evaluator.mIoU()
            acc_tmp = evaluator.acc()
            print(datetime.datetime.now().strftime('[%Y-%m-%d %H:%M:%S]'),
                  ' iter: %-5d    acc: %-10.3f    mIoU: %-10.3f ' % (i , acc_tmp, mIoU_tmp))


if __name__ == '__main__':
    main()